// Copyright 2000-2023 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
package org.jetbrains.kotlin.idea.navigation

import com.intellij.openapi.project.Project
import com.intellij.openapi.roots.GeneratedSourcesFilter
import com.intellij.openapi.vfs.VirtualFile
import com.intellij.psi.PsiElement
import org.jetbrains.kotlin.analysis.api.KaExperimentalApi
import org.jetbrains.kotlin.analysis.api.KaImplementationDetail
import org.jetbrains.kotlin.analysis.api.analyze
import org.jetbrains.kotlin.analysis.api.projectStructure.analysisContextModule
import org.jetbrains.kotlin.psi.KtElement

/**
 * Marks sources generated by Analysis API resolve extensions as generated, and handles
 * substituting in navigation elements according to the logic of the extensions.
 */
class KotlinResolveExtensionGeneratedSourcesFilter : GeneratedSourcesFilter() {
    @OptIn(KaImplementationDetail::class)
    override fun isGeneratedSource(file: VirtualFile, project: Project): Boolean =
        file.analysisContextModule != null

    private val KtElement.hasAnalysisExtensionContext: Boolean
        get() = containingKtFile.virtualFile?.let { isGeneratedSource(it, project) } == true

    @OptIn(KaExperimentalApi::class)
    override fun getOriginalElements(element: PsiElement): List<PsiElement> {
        if (element !is KtElement || !element.hasAnalysisExtensionContext) return emptyList()
        return analyze(element) {
            element.resolveExtensionNavigationElements.toList()
        }
    }
}
